#
# pre-sms framing protocol for binary data longer than maximum size
#
import base64
import math
import string
import zlib


class CorruptMessageError(Exception): pass

class MsgFrag:
    def __init__(self,
                 uniq, crc, seqno, num_frags,
                 bytes, sender_num):
        self.uniq = uniq
        self.crc = crc
        self.seqno = int(seqno)
        self.num_frags = num_frags
        self.bytes = bytes
        self.sender_num = sender_num

class Encoder:
    HDR_TYPE_LEN   = 1
    HDR_CHKSUM_LEN = 1

    HDR_TYPE_SINGLE     = '1'
    HDR_TYPE_MULTI      = 'M'
    HDR_TYPE_MULTI_CONT = 'm'

    MAX_FRAGS = 100
    
    def __init__(self, max_sms_len):
        self.max_sms_len = max_sms_len
        self.max_b64_len = max_sms_len/4*3
        
    @staticmethod
    def b64len(nbytes):
        return math.ceil((len(bytes) + Encoder.HDR_TYPE_LEN +
                          Encoder.HDR_CHKSUM)/3.0)*4

    @staticmethod
    def crc(bytes):
        """El cheapo CRC, constructed from taking CRC32 and shrinking
        it. Probably not very good collision properties."""
        crc32 = zlib.crc32(bytes)
        return ((crc32 & 0xFF) ^ (crc32>>8 & 0xFF) ^
                (crc32>>16 & 0xFF) ^ (crc32>>24 & 0xFF))
        
    def encode(self, uniq, bytes):
        b64len = Encoder.HDR_TYPE_LEN + Encoder.HDR_CHKSUM_LEN + len(bytes)
        if b64len > self.max_b64_len:
            return self.encode_multiple(uniq, bytes)
        else:
            return self.encode_single(bytes)

    def encode_single(self, bytes):
        """Message format: 1 CHKSUM bytes"""
        s = Encoder.HDR_TYPE_SINGLE + chr(Encoder.crc(bytes)) + bytes
        s = base64.b64encode(s)
        assert len(s) <= self.max_sms_len
        return [s]

    def encode_multiple(self, uniq, bytes):
        """First message: M FRAGS CHKSUM UNIQ \0 bytes
           other message: m SEQNO UNIQ \0 bytes"""
        msgs = []
        seqno = 0

        crc = Encoder.crc(bytes)
        l = self.max_b64_len - (Encoder.HDR_TYPE_LEN +
                                Encoder.HDR_CHKSUM_LEN +
                                len(uniq) +
                                1 + # \0
                                1)  # FRAGS
        
        first_msg = ('%s%c%c%s\0' % (Encoder.HDR_TYPE_MULTI, 0, crc, uniq)) + bytes[:l]
        bytes = bytes[l:]
        seqno = seqno + 1
        
        while len(bytes) > 0:
            l = self.max_b64_len - (Encoder.HDR_TYPE_LEN +
                                      len(uniq) +
                                      1 + # \0
                                      1)  # SEQNO
            s = ('%s%c%s\0' % (Encoder.HDR_TYPE_MULTI_CONT, seqno, uniq)) + bytes[:l]
            msgs.append(base64.b64encode(s))
            bytes = bytes[l:]
            seqno = seqno + 1

        # patch in how many fragments
        first_msg = first_msg[0] + chr(seqno) + first_msg[2:]
        msgs = [base64.b64encode(first_msg)] + msgs
        
        for m in msgs:
            assert len(m) < self.max_sms_len
        return msgs

    def decode(self, b64, sender_num = None):
        if not sender_num:
            sender_num = 'XXXXXXXXX'
        
        bytes = base64.b64decode(b64)
        if bytes[0] == Encoder.HDR_TYPE_SINGLE:
            chk = ord(bytes[1])
            return MsgFrag(uniq = None, crc = chk,
                           seqno = 0, num_frags = 1,
                           bytes = bytes[2:],
                           sender_num = sender_num)
        elif bytes[0] == Encoder.HDR_TYPE_MULTI:
            frags = ord(bytes[1])
            if frags > Encoder.MAX_FRAGS:
                raise CorruptMessageError('too many fragments')
            chk = ord(bytes[2])
            bytes = bytes[3:]
            null = bytes.find('\0')
            if null == -1:
                raise CorruptMessageError('unterminated unique string')
            uniq = bytes[:null]
            data = bytes[null + 1:]
            return MsgFrag(uniq = uniq, crc = chk,
                           seqno = 0, num_frags = frags,
                           bytes = data,
                           sender_num = sender_num)
        elif bytes[0] == Encoder.HDR_TYPE_MULTI_CONT:
            seqno = ord(bytes[1])
            bytes = bytes[2:]
            null = bytes.find('\0')
            if null == -1:
                raise CorruptMessageError('unterminated unique string')            
            uniq = bytes[:null]
            data = bytes[null + 1:]
            return MsgFrag(uniq = uniq, crc = None,
                           seqno = seqno, num_frags = None,
                           bytes = data,
                           sender_num = sender_num)
        else:
            raise CorruptMessageError('unrecognized initial byte')

    def merge(self, frags):
        def tup_cmp(a, b):
            return cmp(a.seqno, b.seqno)
        frags.sort(tup_cmp)
        if frags[0].num_frags is None:
            return None
        assert frags[0].num_frags >= len(frags)
        if frags[0].num_frags != len(frags):
            return None
        s = string.join([f.bytes for f in frags], '')

        if Encoder.crc(s) != frags[0].crc:
            raise CorruptMessageError('fails crc check')

        return s

if __name__ == '__main__':
    import random

    def randstr(len):
        s = ''
        for i in range(len):
            s = s + chr(random.randint(0, 255))
        return s
    
    enc = Encoder(159)

    e = {}
    ss = [('AcE', 'abcde'), ('BDF', 'abcdefghijklmnop'), ('cde', 'abcdefg'*20)]

    for l in range(100,1000):
        for j in range(10):
            s = randstr(l)
            u = base64.b64encode(randstr(6))
            ss.append((u,s))
    
    for i in range(len(ss)):
        e[i] = enc.encode(ss[i][0], ss[i][1])
        # print e[i]
    
    for i in range(len(ss)):
	dec = enc.merge([enc.decode(j) for j in e[i]]) 
        assert dec == ss[i][1], 'orig = %s; dec = %s' % (repr(ss[i][1]),
		repr(dec))
